package away3d.animators.data ;


import away3d.animators.skeleton.JointPose ;
import away3d.animators.skeleton.Skeleton ;
import away3d.animators.skeleton.SkeletonJoint ;
import away3d.animators.skeleton.SkeletonPose ;
import away3d.animators.skeleton.SkeletonTreeNode ;
import away3d.core.base.IRenderable ;
import away3d.core.base.SkinnedSubGeometry ;
import away3d.core.base.SubMesh ;
import away3d.core.math.Quaternion ;
import away3d.entities.Mesh ;
import away3d.materials.passes.MaterialPassBase ;
import flash.display3D.Context3D ;
import flash.display3D.Context3DProgramType ;
import flash.geom.Vector3D ;
import flash.utils.Dictionary ;



/**
 * SkeletonAnimationState defines the state for a given Mesh and SkeletonAnimation. The state consists of a skeleton pose.
 *
 * @see away3d.core.animation.skinned.SkinnedAnimation
 *
 */
class SkeletonAnimationState extends AnimationStateBase
{
    
    
    private var _globalMatrices         : Vector<Float> ;
    private var _numJoints              : Int ;
    private var _skinnedAnimation       : SkeletonAnimation ;
    private var _jointsPerVertex        : Int ;
    private var _bufferFormat           : String ;
    private var _skeleton               : Skeleton ;
    private var _blendTree              : SkeletonTreeNode ;
    private var _globalPose             : SkeletonPose ;
    private var _globalInput            : Boolean ;
    private var _buffersValid           : Dictionary< ,Bool > ;//TODO : resolve this 
    private var _globalMatricesInvalid  : Boolean ;
    
    
    
    /**
     * Creates a SkeletonAnimationState object.
     * @param animation The animation object the state refers to.
     * @param jointsPerVertex The amount of joints per vertex define
     */
    public function new( animation : SkeletonAnimation )
    {
        
        super( animation ) ;
        _buffersValid       = new Dictionary( true ) ;
        _skinnedAnimation   = animation ;
        
        if( animation.numJoints > 0 )
        {
            
            init() ;
            
        }
        
    }
    
    
    private function init()
    {
        
        _jointsPerVertex    = _skinnedAnimation.jointsPerVertex ;
        _skeleton           = _skinnedAnimation.skeleton ;
        _numJoints          = _skinnedAnimation.numJoints ;
        _globalMatrices     = new Vector( _numJoints * 12, true ) ;
        _bufferFormat       = "float" + _jointsPerVertex ;
        _globalPose         = new SkeletonPose() ;
        
        var j : Int = 0 ;// TODO: Check if this is correct
        for( i in 0..._numJoints )
        {
            
            _globalMatrices[ j++ ] = 1; _globalMatrices[ j++ ] = 0; _globalMatrices[ j++ ] = 0; _globalMatrices[ j++ ] = 0;
            _globalMatrices[ j++ ] = 0; _globalMatrices[ j++ ] = 1; _globalMatrices[ j++ ] = 0; _globalMatrices[ j++ ] = 0;
            _globalMatrices[ j++ ] = 0; _globalMatrices[ j++ ] = 0; _globalMatrices[ j++ ] = 1; _globalMatrices[ j++ ] = 0;
        
        }
        
    }
    
    
    /**
     * The amount of joints in the target skeleton.
     */
    public function get_numJoints() : Int
    {
        
        return _numJoints ;
        
    }
    
    
    /**
     * The chained raw data of the global pose matrices in row-major order.
     */
    public function get_globalMatrices() : Vector<Float>
    {
        
        return _globalMatrices ;
        
    }
    
    /**
     * The global skeleton pose used to transform the mesh's vertices.
     */
    public function get_globalPose() : SkeletonPose
    {
        
        if( _stateInvalid ) updateGlobalPose() ;
        
        return _globalPose ;
        
    }
    
    
    public function set_globalPose( value : SkeletonPose ): SkeletonPose
    {
        
        if( !_globalInput ) throw new Error( "Cannot set global pose if globalInput is false" ) ;
        
        _globalPose             = value ;
        _globalMatricesInvalid  = true ;
        
        return value ;
        
    }
    
    
    public function validateGlobalMatrices()
    {
        
        _stateInvalid           = false ;
        _globalMatricesInvalid  = false ;
        
    }
    
    
    public var blendTree( get_blendTree, set_blendTree ): SkeletonTreeNode ; 
    
    
    /**
     * The local skeleton blend tree that will be used to generate the global pose.
     */
    public function get_blendTree() : SkeletonTreeNode
    {
        
        return _blendTree ;
        
    }
    
    
    public function set_blendTree( value : SkeletonTreeNode ): SkeletonTreeNode
    {
        
        _blendTree = value ;
        return SkeletonTreeNode ;
        
    }
    
    
    /**
     * @inheritDoc
     */
    override public function setRenderState ( context       : Context3D
                                            , contextIndex  : Int
                                            , pass          : MaterialPassBase
                                            , renderable    : IRenderable
                                            )
    {   
        
        if( _numJoints == 0 )
        {
            // delayed skeleton instantiation
            if( _skinnedAnimation.numJoints > 0 )
            {
                init() ;
            }
            else
            {
                return ;
            }
            
        }
        
        // do on request of globalPose
        if( _stateInvalid ) updateGlobalPose() ;
        if( _globalMatricesInvalid ) convertToMatrices() ;
            
        if( _skinnedAnimation.usesCPU ) 
        {
            
            var subGeom : SkinnedSubGeometry = SkinnedSubGeometry( SubMesh( renderable ).subGeometry ) ;
            
            if( !_buffersValid[ subGeom ] )
            {
                morphGeometry( subGeom ) ;
                _buffersValid[ subGeom ] = true;
            }
            
            return;
            
        }
        
        var skinnedGeom : SkinnedSubGeometry = SkinnedSubGeometry(SubMesh(renderable).subGeometry) ;
        
        context.setProgramConstantsFromVector   (   Context3DProgramType.VERTEX
                                                ,   pass.numUsedVertexConstants
                                                ,   _globalMatrices
                                                ,   _numJoints*3
                                                ) ;
        
        var streamOffset = pass.numUsedStreams ;
        
        context.setVertexBufferAt   ( streamOffset
                                    , skinnedGeom.getJointIndexBuffer( context, contextIndex )
                                    , 0
                                    , _bufferFormat 
                                    ) ;
                                    
        context.setVertexBufferAt   ( streamOffset + 1
                                    , skinnedGeom.getJointWeightsBuffer( context, contextIndex )
                                    , 0
                                    , _bufferFormat 
                                    ) ;
        
    }
    
    
    private function updateGlobalPose()
    {
        
        if ( !_globalInput )
        {
            
            _blendTree.updatePose( _skeleton ) ;
            _blendTree.skeletonPose.toGlobalPose( _globalPose, _skeleton ) ;
                
        }
        
        _globalMatricesInvalid  = true;
        _stateInvalid           = false;
        
    }
    
    
    private function updateGlobalPose()
    {
        
        if( !_globalInput )
        {
            
            _blendTree.updatePose( _skeleton ) ;
            _blendTree.skeletonPose.toGlobalPose( _globalPose, _skeleton ) ;
            
        }
        
        _globalMatricesInvalid  = true ;
        _stateInvalid           = false ;
        
    }
    
    
    /**
     * @inheritDoc
     */ 
    override public function clone() : AnimationStateBase
    {
        
        return new SkeletonAnimationState( _skinnedAnimation ) ;
        
    }
    
    
    public function globalInput( get_globalInput, set_globalInput ): Bool ;
    
    /**
     * Defines whether or not to bypass the blend tree and allow setting the global skeleton pose directly.
     * todo: remove, use post-processing effects for global position based... stuff.
     */
    public function get_globalInput() : Bool
    {
        
        return _globalInput ;
        
    }
    
    
    public function set_globalInput( value: Bool ) : Bool
    {
        
        _globalInput = value ;
        return value ;
        
    }
    
    
    /**
     * Converts the current final pose to matrices for the actual transformations
     */
    private function convertToMatrices()
    {
        
        // convert pose to matrix
        var mtxOffset   : Int;
        var globalPoses                     = _globalPose.jointPoses;
        var raw         : Vector<Float> ;
        var ox          : Float ;
        var oy          : Float ;
        var oz          : Float ;
        var ow          : Float ;
        var xy2         : Float ;
        var xz2         : Float ;
        var xw2         : Float ;
        var yz2         : Float ;
        var yw2         : Float ;
        var zw2         : Float ;
        var xx          : Float ;
        var yy          : Float ;
        var zz          : Float ;
        var ww          : Float ;
        
        var n11         : Float ;
        var n12         : Float ;
        var n13         : Float ;
        var n14         : Float ;
        var n21         : Float ;
        var n22         : Float ;
        var n23         : Float ;
        var n24         : Float ;
        var n31         : Float ;
        var n32         : Float ;
        var n33         : Float ;
        var n34         : Float ;
        var m11         : Float ;
        var m12         : Float ;
        var m13         : Float ;
        var m14         : Float ;
        var m21         : Float ;
        var m22         : Float ;
        var m23         : Float ;
        var m24         : Float ;
        var m31         : Float ;
        var m32         : Float ;
        var m33         : Float ;
        var m34         : Float ;
        var joints                          = _skeleton.joints;
        var pose        : JointPose;
        var quat        : Quaternion;
        var vec         : Vector3D;
        
        for( i in  0..._numJoints ) 
        {
            
            pose    = globalPoses[ i ] ;
            quat    = pose.orientation ;
            vec     = pose.translation ;
            
            ox      = quat.x ;
            oy      = quat.y ;
            oz      = quat.z ;
            ow      = quat.w ;
            
            
            xy2     = 2.0 * ox * oy ;
            xz2     = 2.0 * ox * oz ;
            xw2     = 2.0 * ox * ow ;
            
            yz2     = 2.0 * oy * oz ;
            yw2     = 2.0 * oy * ow ;
            zw2     = 2.0 * oz * ow ;
            
            xx      = ox * ox ;
            yy      = oy * oy ; 
            zz      = oz * oz ;
            ww      = ow * ow ;
            
            n11     = xx - yy - zz + ww ;
            n12     = xy2 - zw2 ;
            n13     = xz2 + yw2 ;
            n14     = vec.x ;
            
            n21     = xy2 + zw2 ;
            n22     = -xx + yy - zz + ww ; 
            n23     = yz2 - xw2 ;
            n24     = vec.y ;
            
            n31     = xz2 - yw2 ;
            n32     = yz2 + xw2 ;
            n33     = -xx - yy + zz + ww ;
            n34     = vec.z ;
            
            // prepend inverse bind pose
            raw     = joints[ i ].inverseBindPose ;
            
            m11 = raw[ 0 ] ;    m12 = raw[ 4 ] ;    m13 = raw[ 8 ] ;     m14 = raw[ 12 ] ;
            m21 = raw[ 1 ] ;    m22 = raw[ 5 ] ;    m23 = raw[ 9 ] ;     m24 = raw[ 13 ] ;
            m31 = raw[ 2 ] ;    m32 = raw[ 6 ] ;    m33 = raw[ 10 ] ;    m34 = raw[ 14 ] ;
            
            _globalMatrices[ mtxOffset++ ] = n11 * m11 + n12 * m21 + n13 * m31 ;
            _globalMatrices[ mtxOffset++ ] = n11 * m12 + n12 * m22 + n13 * m32 ;
            _globalMatrices[ mtxOffset++ ] = n11 * m13 + n12 * m23 + n13 * m33 ;
            _globalMatrices[ mtxOffset++ ] = n11 * m14 + n12 * m24 + n13 * m34 + n14 ;
            _globalMatrices[ mtxOffset++ ] = n21 * m11 + n22 * m21 + n23 * m31 ;
            _globalMatrices[ mtxOffset++ ] = n21 * m12 + n22 * m22 + n23 * m32 ;
            _globalMatrices[ mtxOffset++ ] = n21 * m13 + n22 * m23 + n23 * m33 ;
            _globalMatrices[ mtxOffset++ ] = n21 * m14 + n22 * m24 + n23 * m34 + n24 ;
            _globalMatrices[ mtxOffset++ ] = n31 * m11 + n32 * m21 + n33 * m31 ;
            _globalMatrices[ mtxOffset++ ] = n31 * m12 + n32 * m22 + n33 * m32 ;
            _globalMatrices[ mtxOffset++ ] = n31 * m13 + n32 * m23 + n33 * m33 ;
            _globalMatrices[ mtxOffset++ ] = n31 * m14 + n32 * m24 + n33 * m34 + n34 ;
            
        }
        
        _globalMatricesInvalid = false ;
        
    }
    
    
    /**
     * If the animation can't be performed on cpu, transform vertices manually
     * @param subGeom The subgeometry containing the weights and joint index data per vertex.
     * @param pass The material pass for which we need to transform the vertices
     *
     * todo: we may be able to transform tangents more easily, similar to how it happens on gpu
     */
    private function morphGeometry( subGeom: SkinnedSubGeometry )
    {
        
        var verts               = subGeom.vertexData ;
        var normals             = subGeom.vertexNormalData ;
        var tangents            = subGeom.vertexTangentData ;
        var targetVerts         = subGeom.animatedVertexData ;
        var targetNormals       = subGeom.animatedNormalData ;
        var targetTangents      = subGeom.animatedTangentData ;
        var jointIndices        = subGeom.jointIndexData ;
        var jointWeights        = subGeom.jointWeightsData ;
        
        var i1          : Int ; 
        var i2                  = 1 ; 
        var i3                  = 2 ;
        var j           : Int ;
        var k           : Int ;
        var vx          : Float ;
        var vy          : Float ;
        var vz          : Float ;
        var nx          : Float ;
        var ny          : Float ;
        var nz          : Float ;
        var tx          : Float ;
        var ty          : Float ;
        var tz          : Float ;
        var len                 = verts.length;
        var weight      : Float ;
        var mtxOffset   : Int ;
        
        var vertX       : Float ;
        var vertY       : Float ;
        var vertZ       : Float ;
        
        var normX       : Float ;
        var normY       : Float ;
        var normZ       : Float ;
        
        var tangX       : Float ;
        var tangY       : Float ;
        var tangZ       : Float ;
        
        var m11         : Float ;
        var m12         : Float ;
        var m13         : Float ;
        
        var m21         : Float ;
        var m22         : Float ;
        var m23         : Float ;
        var m31         : Float ;
        var m32         : Float ;
        var m33         : Float ;
        
        while( i1 < len )
        {
            
            vertX   = verts[ i1 ] ;
            vertY   = verts[ i2 ] ;
            vertZ   = verts[ i3 ] ;
            
            vx      = 0 ;
            vy      = 0 ;
            vz      = 0 ;
            
            normX   = normals[ i1 ] ;
            normY   = normals[ i2 ] ;
            normZ   = normals[ i3 ] ;
            
            nx      = 0 ;
            ny      = 0 ;
            nz      = 0 ;
            
            tangX   = tangents[ i1 ] ;
            tangY   = tangents[ i2 ] ;
            tangZ   = tangents[ i3 ] ;
            
            tx      = 0 ; 
            ty      = 0 ;
            tz      = 0 ;
            
            k       = 0;
            
            while( k < _jointsPerVertex )
            {
                
                weight = jointWeights[ j ] ;
                
                if( weight == 0 )
                {
                    
                    j += _jointsPerVertex - k ;
                    k = _jointsPerVertex ;
                    
                }
                else 
                {
                    
// implicit /3*12 (/3 because indices are multiplied by 3 for gpu matrix access, *12 because it's the matrix size)
                    mtxOffset = jointIndices[ j++ ]*4 ;
                    
                    m11 = _globalMatrices[ mtxOffset ] ; 
                    m12 = _globalMatrices[ mtxOffset + 1 ] ; 
                    m13 = _globalMatrices[ mtxOffset + 2 ] ;
                    
                    m21 = _globalMatrices[ mtxOffset + 4 ] ; 
                    m22 = _globalMatrices[ mtxOffset + 5 ] ; 
                    m23 = _globalMatrices[ mtxOffset + 6 ] ;
                    
                    m31 = _globalMatrices[ mtxOffset + 8 ] ; 
                    m32 = _globalMatrices[ mtxOffset + 9 ] ; 
                    m33 = _globalMatrices[ mtxOffset + 10 ] ;
                    
                    vx += weight*( m11*vertX + m12*vertY + m13*vertZ + _globalMatrices[ mtxOffset + 3 ] ) ;
                    vy += weight*( m21*vertX + m22*vertY + m23*vertZ + _globalMatrices[ mtxOffset + 7 ] ) ;
                    vz += weight*( m31*vertX + m32*vertY + m33*vertZ + _globalMatrices[ mtxOffset + 11 ] ) ;
                    
                    nx += weight*( m11*normX + m12*normY + m13*normZ ) ;
                    ny += weight*( m21*normX + m22*normY + m23*normZ ) ;
                    nz += weight*( m31*normX + m32*normY + m33*normZ ) ;
                    tx += weight*( m11*tangX + m12*tangY + m13*tangZ ) ;
                    ty += weight*( m21*tangX + m22*tangY + m23*tangZ ) ;
                    tz += weight*( m31*tangX + m32*tangY + m33*tangZ ) ;
                    k++;
                
                }
            }
            
            targetVerts[ i1 ]       = vx ; 
            targetVerts[ i2 ]       = vy ;
            targetVerts[ i3 ]       = vz ;
            
            targetNormals[ i1 ]     = nx ;
            targetNormals[ i2 ]     = ny ;
            targetNormals[ i3 ]     = nz ;
            
            targetTangents[ i1 ]    = tx ;
            targetTangents[ i2 ]    = ty ;
            targetTangents[ i3 ]    = tz ;
            
            i1                      += 3 ; 
            i2                      += 3 ; 
            i3                      += 3 ;
            
        }
        
        subGeom.animatedVertexData  = targetVerts ;
        subGeom.animatedNormalData  = targetNormals ;
        subGeom.animatedTangentData = targetTangents ;
        
    }
    
    
    public function applyRootDelta()
    {
        
        var delta       = blendTree.rootDelta ;
        var dist        = delta.length ;
        var len : Int ;
        
        if( dist > 0 )
        {
            len = _owners.length;
            for( i in 0...len )
            {
                
                _owners[ i ].translateLocal( delta, dist );
                
            }
            
        }
        
    }
    
    
}